package com.github.maojx0630.onebot.socket;

import java.util.Map;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

/**
 * @author 毛家兴
 * @since 2024-04-02 09:59
 */
@Slf4j
@AllArgsConstructor
public class WsHandshakeInterceptor implements HandshakeInterceptor {

  private WsTokenService wsTokenService;

  @Override
  public boolean beforeHandshake(
      ServerHttpRequest request,
      ServerHttpResponse response,
      WebSocketHandler wsHandler,
      Map<String, Object> attributes) {
    WsAuth wsAuth = new WsAuth(request);
    if (wsAuth.isAuth()) {
      if (wsTokenService.checkToken(wsAuth.getQq(), wsAuth.getToken())) {
        attributes.put("qq", wsAuth.getQq());
        attributes.put("token", wsAuth.getToken());
        return true;
      }
    }
    log.info("握手失败,未获取到正确的qq号码或token!---:{}", wsAuth);
    return false;
  }

  @Override
  public void afterHandshake(
      ServerHttpRequest request,
      ServerHttpResponse response,
      WebSocketHandler wsHandler,
      Exception exception) {}
}
